#include "stdafx.h"

#include <stdio.h>
#include <tchar.h>
#include <string.h>
#include <ws2spi.h>
#include <sporder.h>
#pragma comment(lib,"ws2_32.lib")

#include "SPITool.h"
#include "SPIGlobal.h"

int totalprotos=0;
DWORD protoinfosize=0;
LPWSAPROTOCOL_INFOW protoinfo=NULL;

BOOL getfilter()
{
	int errorcode;

	protoinfo=NULL;
	totalprotos=0;
	protoinfosize=0;

	if(WSCEnumProtocols(NULL,protoinfo,&protoinfosize,&errorcode)==SOCKET_ERROR)
	{
		if(errorcode!=WSAENOBUFS)
		{
			_tprintf(_T("First WSCEnumProtocols Error: %d\n"),errorcode);
			return FALSE;
		}
	}

	if((protoinfo=(LPWSAPROTOCOL_INFOW)GlobalAlloc(GPTR,protoinfosize))==NULL)
	{
		_tprintf(_T("GlobalAlloc in getfilter Error: %d\n"),GetLastError());
		return FALSE;
	}

	if((totalprotos=WSCEnumProtocols(NULL,protoinfo,&protoinfosize,&errorcode))==SOCKET_ERROR)
	{
		_tprintf(_T("Second WSCEnumProtocols Error: %d\n"),GetLastError());
		return FALSE;
	}

	_tprintf(_T("Found %d protocols!\n"),totalprotos); 
	return TRUE;
}

void freefilter()
{
	GlobalFree(protoinfo);
}

bool installfilter(const TCHAR* pFilePath)
{
	int i;
	int provcnt;
	int cataindex;
	int errorcode;
	BOOL rawip=FALSE;
	BOOL udpip=FALSE;
	DWORD iplayercataid=0,udporigcataid=0;
	TCHAR filter_path[MAX_PATH]; 
	TCHAR filter_name[MAX_PATH];
	TCHAR chainname[WSAPROTOCOL_LEN+1]; 
	LPDWORD cataentries;
	WSAPROTOCOL_INFOW iplayerinfo,udpchaininfo,chainarray[1];

	getfilter();

	for(i=0;i<totalprotos;i++)
	{
		if(!rawip
			&& protoinfo[i].iAddressFamily==AF_INET
			&& protoinfo[i].iProtocol==IPPROTO_IP)
		{
			rawip=TRUE;
			memcpy(&iplayerinfo,&protoinfo[i],sizeof(WSAPROTOCOL_INFOW));
			iplayerinfo.dwServiceFlags1=protoinfo[i].dwServiceFlags1 & (~XP1_IFS_HANDLES);
		}

		if(!udpip
			&& protoinfo[i].iAddressFamily==AF_INET
			&& protoinfo[i].iProtocol==IPPROTO_TCP)
		{
			udpip=TRUE;
			udporigcataid=protoinfo[i].dwCatalogEntryId;
			memcpy(&udpchaininfo,&protoinfo[i],sizeof(WSAPROTOCOL_INFOW));
			udpchaininfo.dwServiceFlags1=protoinfo[i].dwServiceFlags1 & (~XP1_IFS_HANDLES);
		}
	}

	_tcscpy(iplayerinfo.szProtocol,_T("NetEast-IpFilter"));
	iplayerinfo.ProtocolChain.ChainLen=LAYERED_PROTOCOL;

	/*if(GetCurrentDirectory(MAX_PATH,filter_path)==0)
	{
		_tprintf(_T("GetCurrentDirectory Error: %d\n"),GetLastError());
		return false;
	}*/
	//_tcscpy(filter_name,_T("\\ipfilter.dll")); 
	//_tcscat(filter_path,filter_name);

	_tcscpy(filter_path,pFilePath); //dll path

	if(WSCInstallProvider(&filterguid,filter_path,&iplayerinfo,1,&errorcode)==
		SOCKET_ERROR)
	{
		_tprintf(_T("WSCInstallProvider Error: %d\n"),errorcode);
		return false;
	}

	freefilter();

	getfilter();

	for(i=0;i<totalprotos;i++)
	{
		if(memcmp(&protoinfo[i].ProviderId,&filterguid,sizeof(GUID))==0)
		{
			iplayercataid=protoinfo[i].dwCatalogEntryId;
			break;
		}
	}

	provcnt=0;
	if(udpip)
	{
		_tcscpy(chainname,_T("NetEast-TcpFilter"));
		_tcscpy(udpchaininfo.szProtocol,chainname);

		if(udpchaininfo.ProtocolChain.ChainLen==BASE_PROTOCOL)
		{
			udpchaininfo.ProtocolChain.ChainEntries[1]=udporigcataid;
		}
		else
		{
			for(i=udpchaininfo.ProtocolChain.ChainLen;i>0;i--)
			{
				udpchaininfo.ProtocolChain.ChainEntries[i+1]=udpchaininfo.ProtocolChain.ChainEntries[i];
			}
		}

		udpchaininfo.ProtocolChain.ChainLen++;
		udpchaininfo.ProtocolChain.ChainEntries[0]=iplayercataid;

		memcpy(&chainarray[provcnt++],&udpchaininfo,sizeof(WSAPROTOCOL_INFOW));
	}

	if(WSCInstallProvider(&filterchainguid,filter_path,chainarray,provcnt,&errorcode)==
		SOCKET_ERROR)
	{
		_tprintf(_T("WSCInstallProvider for chain Error: %d\n"),errorcode);
		return false;
	}

	freefilter();

	getfilter();

	if((cataentries=(LPDWORD)GlobalAlloc(GPTR,totalprotos*sizeof(WSAPROTOCOL_INFOW)))
		==NULL)
	{
		_tprintf(_T("GlobalAlloc int installfilter Error: %d\n"),errorcode);
		return false;
	}

	cataindex=0;
	for(i=0;i<totalprotos;i++)
	{
		if(memcmp(&protoinfo[i].ProviderId,&filterguid,sizeof(GUID))==0 
			|| memcmp(&protoinfo[i].ProviderId,&filterchainguid,sizeof(GUID))==0)
		{
			cataentries[cataindex++]=protoinfo[i].dwCatalogEntryId;
		}
	}

	for(i=0;i<totalprotos;i++)
	{
		if(memcmp(&protoinfo[i].ProviderId,&filterguid,sizeof(GUID))!=0 
			&& memcmp(&protoinfo[i].ProviderId,&filterchainguid,sizeof(GUID))!=0)
		{
			cataentries[cataindex++]=protoinfo[i].dwCatalogEntryId;
		}
	}

	if((errorcode==WSCWriteProviderOrder(cataentries,totalprotos))!=ERROR_SUCCESS)
	{
		_tprintf(_T("WSCWriteProviderOrder Error: %d\n"),GetLastError());
		return false;
	}

	freefilter();

	_tprintf(_T("\nInstall IP Filter Successfully"));
	return true;
}

bool removefilter(const TCHAR* pFilePath)
{
	int errorcode;
	BOOL signal=TRUE;

	if(WSCDeinstallProvider(&filterguid,&errorcode)==SOCKET_ERROR)
	{
		_tprintf(_T("WSCDeinstall filterguid Error: %d\n"),errorcode);
		signal=FALSE;
	}

	if(WSCDeinstallProvider(&filterchainguid,&errorcode)==SOCKET_ERROR)
	{
		_tprintf(_T("WSCDeinstall filterchainguid Error: %d\n"),errorcode);
		signal=FALSE;
	}



	HANDLE hStopEvent = NULL;
	
	do 
	{
		hStopEvent=::OpenEvent(EVENT_MODIFY_STATE, FALSE, EVENT_NAME);
		if (hStopEvent)
		{
			BOOL bRet = ::SetEvent(hStopEvent);
			if (bRet)
			{
				_tprintf(_T("SetEvent ..\n"));
				::CloseHandle(hStopEvent);
			}
			
			
		}


	} while (hStopEvent);
	

	if(signal)
	{
		_tprintf(_T("Deinstall IP Filter Successfully"));
	}
	return true;
}
